"""
UNet推理脚本
对单张图像进行推理并保存结果
"""
import os
import sys
import argparse
import torch

# 添加项目根目录到路径
sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))

from baselines.unet.model import UNet
from baselines.unet.config import MODEL_CONFIG, DATA_CONFIG
from utils.infer_utils import inference_single_image, load_model_for_inference


def main():
    parser = argparse.ArgumentParser(description='UNet推理脚本')
    parser.add_argument('--image', type=str, required=True,
                        help='输入图像路径')
    parser.add_argument('--checkpoint', type=str, required=True,
                        help='模型权重路径')
    parser.add_argument('--output_dir', type=str, default='outputs/unet',
                        help='输出目录')
    parser.add_argument('--threshold', type=float, default=0.5,
                        help='二值化阈值')
    parser.add_argument('--device', type=str, default='cuda',
                        help='设备 (cuda/cpu)')

    args = parser.parse_args()

    # 检查输入文件
    if not os.path.exists(args.image):
        raise FileNotFoundError(f"输入图像不存在: {args.image}")

    if not os.path.exists(args.checkpoint):
        raise FileNotFoundError(f"模型权重不存在: {args.checkpoint}")

    # 设置设备
    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    print(f'[INFO] 使用设备: {device}')

    # 创建模型
    print('[INFO] 创建模型...')
    model = UNet(
        in_channels=MODEL_CONFIG['in_channels'],
        out_channels=MODEL_CONFIG['out_channels'],
        features=MODEL_CONFIG['features']
    )

    # 加载权重
    model = load_model_for_inference(model, args.checkpoint, device)

    # 构建输出路径
    image_name = os.path.basename(args.image)
    image_name_no_ext = os.path.splitext(image_name)[0]
    output_path = os.path.join(args.output_dir, f'{image_name_no_ext}_pred.png')

    # 推理
    print(f'[INFO] 开始推理...')
    print(f'[INFO] 输入图像: {args.image}')
    print(f'[INFO] 模型权重: {args.checkpoint}')

    inference_time, foreground_pixels = inference_single_image(
        model=model,
        image_path=args.image,
        output_path=output_path,
        device=device,
        to_rgb=DATA_CONFIG['to_rgb'],
        threshold=args.threshold
    )

    # 打印信息
    print(f'[INFO] 推理完成！')
    print(f'[INFO] 推理用时: {inference_time:.4f} 秒')
    print(f'[INFO] 前景像素数: {foreground_pixels}')
    print(f'[INFO] 输出路径: {output_path}')


if __name__ == '__main__':
    main()
